# Time-stamp: <2008-03-26 15:41:09 Tao Liu>

"""Module for RangeI, TrackI, FWTrackI and TabFile class

Copyright (c) 2007 Tao Liu <taoliu@jimmy.harvard.edu>

This code is free software; you can redistribute it and/or modify it
under the terms of the BSD License (see the file COPYING included with
the distribution).

@status:  experimental
@version: $Revision$
@author:  Tao Liu
@contact: taoliu@jimmy.harvard.edu
"""

# ------------------------------------
# python modules
# ------------------------------------

# ------------------------------------
# constants
# ------------------------------------
__version__ = "TabIO $Revision$"
__author__ = "Tao Liu <taoliu@jimmy.harvard.edu>"
__doc__ = "RangeI, TrackI, FWTrackI and TabFile classes"

prog_chars = ["-","\\","|","/"]
prog_num = 4
LOG = False
# ------------------------------------
# Misc functions
# ------------------------------------

# ------------------------------------
# Classes
# ------------------------------------

class RangeI:
    """RangeI

    Example:
       >>> reg = RangeI(0,100,"CDK1",-1)
    """
    def __init__ (self, start=0, end=0, name="", strand=0):
        """The coordinate of the reg range is 0-indexed and right-opened.

        strand 1 means plus strand, -1 negative strand.
        """
        self.start = start
        self.end = end
        self.name = name
        self.strand = strand
        self.set = []
        self.length = end-start
        
    def __str__ (self):
        s= "%s:%d..%d" % (str(self.strand),self.start,self.end)
        return s
        #if len(self.set) >0 :
        #    s+=str(len(self.set))
        #    #for i in self.set:
        #    #    s+="("+str(i)+")"

    def intersect(self, range2):
        """Compute the intersect of two RangeI objects. Return the intersect part as RangeI object or False.
        """
        if (self.start<range2.end and self.end >range2.start):
            ordered_coords = [self.start,self.end,range2.start,range2.end]
            ordered_coords.sort()
            return RangeI(
                start=ordered_coords[1],
                end=ordered_coords[2]
                )
                          
        else:
            return False

    def union (self, range2):
        """Compute the union of two RangeI objects. Return the union as RangeI object or False.
        """
        if (self.start<range2.end and self.end >range2.start):
            ordered_coords = [self.start,self.end,range2.start,range2.end]
            ordered_coords.sort()
            return RangeI(
                start=ordered_coords[0],
                end=ordered_coords[3]
                )
        else:
            return False

    def include (self, range2):
        """Test if self is larger than and including range2.
        """
        if (self.start<=range2.start and range2.end <= self.end):
            return True
        else:
            return False


class TrackI:
    """Ranges along the whole genome (commonly with the same
    annotation type), which are stored in a dict.

    Ranges are stored and organized by sequence names (chr names) in a
    dict. They can be sorted by calling self.sort() function.

    Example:
       >>> tabfile = TabFile("tag.bed",format="bed",mode="r")
       >>> track = TrackI()
       >>> for (chromosome,rg) in tabfile:
       ...    range_list.add_range(chromosome,rg)
       >>> track.get_ranges_by_chr["chr1"] # all ranges in chr1 
    
    """
    def __init__ (self,anno=""):
        """The coordnate of the extended region is 1-indexed.

        """
        self.__ranges = {}
        self.__well_sorted = False
        self.__well_merged = False
        self.annotation = anno
        self.total = 0

    def add_range (self, chromosome, reg):
        """Add a range to the list according to the sequence name.
        
        chromosome -- mostly the chromosome name
        reg -- a Range object
        """
        if not isinstance(reg,RangeI):
            raise Exception("Must add a RangeI object!")
        if not self.__ranges.has_key(chromosome):
            self.__ranges[chromosome] = [] # a tuple for (start_pos,end_pos,strand,meta_data)
        self.__ranges[chromosome].append((reg.start,reg.end,reg.strand,None))
        self.total+=1
        self.__well_merged = False
        self.__well_sorted = False

    def get_ranges_by_chr (self, chromosome):
        """Return array of locations by chromosome.

        Not recommanded! Use generate_rangeI_by_chr() instead.
        """
        if self.__ranges.has_key(chromosome):
            return self.__ranges[chromosome]
        else:
            raise Exception("No such chromosome name (%s) in TrackI object!\n" % (chromosome))

    def generate_rangeI_by_chr (self, chromosome):
        """A generator to return RangeI objects on given chromosome. 

        Recommanded function.
        """
        if self.__ranges.has_key(chromosome):
            for i in self.__ranges[chromosome]:
                yield RangeI(start=i[0],end=i[1],strand=i[2])
        else:
            raise Exception("No such chromosome name (%s) in TrackI object!\n" % (chromosome))

    def get_chr_names (self):
        """Return all the chromosome names stored in this track object.
        """
        l = self.__ranges.keys()
        l.sort()
        return l

    def include(self, frs2, verbose=False):
        """Calculate how many ranges in frs2(a TrackI object) are covered by self. self and frs2 must each be sorted and got rid of overlaps by merge_overlap() function.

        If verbose is set to True, return a tuple (number_of_hits, detail_intersection_info),
        else only return the number.
        """
        if self.__well_merged == False or frs2.status()[1] == False:
            return False
        total_num = 0
        self_chrs = self.get_chr_names()
        frs2_chrs = frs2.get_chr_names()
        v=[]
        for k in self_chrs:
            if not frs2_chrs.count(k):
                continue
            self_rl_k = self.generate_rangeI_by_chr(k) 
            frs2_rl_k = frs2.generate_rangeI_by_chr(k)
            rself = self_rl_k.next()
            if verbose: t = k+":"+str(rself)
            rfrs2 = frs2_rl_k.next()
            try:
                while (True):
                    if rself.end < rfrs2.end:
                        # rself not include rfrs2
                        rself = self_rl_k.next()
                        if verbose:
                            if t.find("\t") != -1: 
                                v.append(t)
                            t=k+":"+str(rself)
                        #print "self"+str(rself)
                    else:       # here rself.end must > rfrs2.end
                        if rself.start <= rfrs2.start: 
                            # rself include rfrs2
                            total_num += 1
                            if verbose: t+="\t"+str(rfrs2)
                        rfrs2 = frs2_rl_k.next()
                        #print "frs2"+str(rfrs2)
            except StopIteration:
                if verbose:
                    if t.find("\t") != -1: 
                        v.append(t)
                continue
        if verbose:
            return (total_num,v)
        else:
            return total_num

    def length (self):
        t = 0
        for key in self.__ranges.keys():
            for r in self.__ranges[key]:
                t+= r[1]-r[0]
        return t
        
    def merge_overlap (self):
       """merge overlapped ranges.

       """
       if not self.__well_sorted:
           self.sort()
       self.total = 0
       for k in self.__ranges.keys(): # for each chromosome
            p = 0
            while True:
                try:
                    if (self.__ranges[k][p][0]<self.__ranges[k][p+1][1] and self.__ranges[k][p][1] >self.__ranges[k][p+1][0]):
                        ordered_coords = [self.__ranges[k][p][0],self.__ranges[k][p+1][1],self.__ranges[k][p][1],self.__ranges[k][p+1][0]]
                        ordered_coords.sort()
                        
                        self.__ranges[k][p]=(ordered_coords[0],ordered_coords[1],1)
                        self.__ranges[k].pop(p+1)
                    else:
                        p+=1
                except IndexError:
                    break
            if (len(self.__ranges[k]) != (p+1)): raise Exception("%d vs %d\n" % (len(self.__ranges[k]),(p+1)))
            self.total += (p+1)
       self.__well_merged = True
       return True

    def status (self):
        """Return the information whether or not the track is sorted and/or merged.

        Return value: tuple (boolean,boolean) means (sorted?,merged?)
        """
        return (self.__well_sorted,self.__well_merged)

    def sort (self):
        if self.__well_sorted:
            pass
        for key in self.__ranges.keys():
            self.__ranges[key].sort(lambda x,y: cmp(x[0],y[0]))
        self.__well_sorted = True

    def __str__ (self):
        return self.__to_wiggle()
        
    def __to_wiggle (self):
        t = "track type=wiggle_0 name=\"Bed Format\" description=\"%s\"\n" % (self.annotation)
        for key in self.__ranges.keys():
            for i in self.__ranges[key]:
                t += "%s\t%d\t%d\n" % (key,i.start,i.end)
        return t


class FWTrackI:
    """Fixed Width Ranges along the whole genome (commonly with the
    same annotation type), which are stored in a dict.

    Ranges are stored and organized by sequence names (chr names) in a
    dict. They can be sorted by calling self.sort() function.

    Example:
       >>> tabfile = TabFile("tag.bed",format="bed",mode="r")
       >>> track = FWTrackI()
       >>> for (chromosome,rg) in tabfile:
       ...    track.add_range(chromosome,rg)
       >>> track.get_ranges_by_chr["chr1"] # all ranges in chr1 
    """
    def __init__ (self,fw=0,anno=""):
        """fw is the fixed-width for all ranges
        """
        self.fw = fw
        self.__ranges = {}
        self.__well_sorted = False
        self.__well_merged = False
        self.total = 0
        self.annotation = anno   # need to be figured out

    def add_range (self, chromosome, range):
        """Add a range to the list according to the sequence name.
        
        chromosome -- mostly the chromosome name
        range -- RangeI objectleverage
        """
        if not isinstance(range,RangeI):
            raise Exception("Must add a RangeI object!")
        if not self.__ranges.has_key(chromosome):
            self.__ranges[chromosome] = []
        #print "add: "+str(range)
        self.__ranges[chromosome].append(range.strand*range.start)
        self.total+=1
        self.__well_sorted = False
        self.__well_merged = False

        
    def get_ranges_by_chr (self, chromosome):
        """Return array of locations by chromosome.

        Not recommanded! Use generate_rangeI_by_chr() instead.
        """
        if self.__ranges.has_key(chromosome):
            return self.__ranges[chromosome]
        else:
            raise Exception("No such chromosome name (%s) in TrackI object!\n" % (chromosome))

    def generate_rangeI_by_chr (self, chromosome):
        """A generator to return RangeI objects on given chromosome. 

        Recommanded function.
        """
        if self.__ranges.has_key(chromosome):
            for i in self.__ranges[chromosome]:
                if i>0:
                    yield RangeI(start=i,end=i+self.fw,strand=1)
                else:
                    yield RangeI(start=-1*i-self.fw,end=-1*i,strand=-1)
        else:
            raise Exception("No such chromosome name (%s) in TrackI object!\n" % (chromosome))

    def get_chr_names (self):
        """Return all the chromosome names stored in this track object.
        """
        l = self.__ranges.keys()
        l.sort()
        return l

    def include(self, frs2, verbose=False):
        """Calculate how many ranges in frs2(a TrackI object) are covered by self.self and frs2 must each be sorted and got rid of overlaps by merge_overlap() function.

        return the number.
        """
        if self.__well_merged == False or frs2.status()[1] == False:
            return False
        total_num = 0
        self_chrs = self.get_chr_names()
        frs2_chrs = frs2.get_chr_names()
        v = []
        for k in self_chrs:
            if not frs2_chrs.count(k):
                continue
            self_rl_k = self.generate_rangeI_by_chr(k) 
            frs2_rl_k = frs2.generate_rangeI_by_chr(k)
            pself = 0
            pfrs2 = 0
            rself = self_rl_k.next() 
            if verbose: t = k+":"+str(rself)
            rfrs2 = frs2_rl_k.next() 
            try:
                while (True):
                    if rself.end < rfrs2.end:
                        # rself not include rfrs2
                        rself = self_rl_k.next()
                        if verbose:
                            if t.find("\t") != -1: 
                                v.append(t)
                            t = k+":"+str(rself)
                    else:       # here rself.end must > rfrs2.end
                        if rself.start <= rfrs2.start: 
                            # rself include rfrs2
                            total_num += 1
                            if verbose: t += "\t"+str(rfrs2)
                        rfrs2 = frs2_rl_k.next()
            except StopIteration:
                if verbose:
                    if t.find("\t") != -1: 
                        v.append(t)
                continue
        if verbose: 
            return (total_num,v)
        else:
            return total_num

    def length (self):
        return self.total*self.fw

    def merge_overlap (self):
       """merge the SAME ranges.
*Note: different with the merge_overlap() in TrackI class, which merges the overlapped ranges. 

       """
       if not self.__well_sorted:
           self.sort()
       self.total = 0
       for k in self.__ranges.keys(): # for each chromosome
           n = 1
           while True:
               try:
                   if self.__ranges[k][n-1]==self.__ranges[k][n]:
                       del(self.__ranges[k][n])
                   else:
                       n += 1
               except IndexError:
                   break
           self.total +=  len(self.__ranges[k])
       self.__well_merged = True

    def status (self):
        """Return the information whether or not the track is sorted and/or merged.

        Return value: tuple (boolean,boolean) means (sorted?,merged?)
        """
        return (self.__well_sorted,self.__well_merged)

    def sort (self):
        """Sort all ranges in list. Recommanded after you import all ranges.
        """
        if self.__well_sorted:
            pass
        for key in self.__ranges.keys():
            self.__ranges[key].sort(lambda x,y: cmp(abs(x),abs(y)))
        self.__well_sorted = True


    def __str__ (self):
        return self.__to_wiggle()
        
    def __to_wiggle (self):
        t = "track type=wiggle_0 name=\"variableStep\" description=\"%s\"\n" % (self.annotation)
        for key in self.__ranges.keys():
            t += "variableStep chrom=%s span=%d\n" % (key,self.fw)
            for i in self.__ranges[key]:
                t += str(abs(i)+1)+"\n"
        return t

    def too_bad (self):
        t=""
        for key in self.__ranges.keys():
            for i in self.__ranges[key]:
                t += key+"\t"+str(abs(i)-10)+"\t"+str(abs(i)+10)+"\n"
        return t
        

class TabFile (file):
    """File Class for tabular File.

    example:
       >>> tabfile = TabFile("tag.bed",format="bed",mode="r")
       >>> range_list = tabfile.build_track()

    If your Tabular File is not in "BED" or "GFF" format, please give a
    particular format string.
    """
    def __init__ (self, name, format="bed", mode="r", buffering=0):
        """All parameters except for format are the same as python's
        builtin file class.

        Format parameter can be "bed","gff" or other particular string
        for other format of file. For example, if your file is '\\t'
        delimited, and the first column of the file is chromosome
        name, the second column is start position, third column is end
        position, fourth column is the strand, the coordinates are
        0-indexed and the range is closed, then you should write the
        format string as "123401\\t" (six numbers and the delimiter).

        Note: Use the fifth and sixth number in format string to
        indicate whether the coordinates are 0-indexed (0) or
        1-indexed (1) and whether the range is opened (0) or closed
        (1) (i.e. whether the end position is included in the range
        or not)
        """
        file.__init__( self, name, mode, buffering)
        if format == "bed" or format == "BED":
            self.__format__ = "123600\t"
        elif format == "gff" or format == "GFF":
            self.__format__ =  "145711\t"
        elif re.match("\d\d\d\d.*",format):
            self.__format__ = format
        else:
            raise self.FormatStringError(format)
        
        self.__delimiter__ = self.__format__[6:]
        self.__chr_col__ = int(self.__format__[0])-1
        self.__start_col__ = int(self.__format__[1])-1
        self.__end_col__ = int(self.__format__[2])-1
        self.__strand_col__ = int(self.__format__[3])-1
        self.__start_shift__ = int(self.__format__[4])
        self.__end_shift__ = int(self.__format__[5])

    def build_track (self, both_strand=False, do_merge=True):
        """Build TrackI from all lines, return a TrackI object.

        Note: All ranges will be merged (exclude the overlapped
        range) then sorted after the track is built.

        If both_strand is True, it will store strand information in
        TrackI object.

        if do_merge is False, it will not merge the same range after
        the track is built.
        """
        track = TrackI()
        if LOG: sys.stdout.write(" adding ranges...\n")
        i = 0
        m = 0
        if LOG: sys.stdout.write("  "+prog_chars[i])
        for (chromosome,rg) in self:
            if not rg or not chromosome:
                continue
            if not both_strand:
                rg.strand = 1
            track.add_range(chromosome,rg)
            if LOG:
                i+=1
                if i == 50000:
                    m+=1
                    sys.stdout.write("\b"+str(m*50000)+"\n   ")
                    i=0
                sys.stdout.write("\b"+prog_chars[i-i/prog_num*prog_num])
            del rg
        if do_merge:
            if LOG: sys.stdout.write( "\b\bsorting then merging...\n")
            track.merge_overlap()
        if LOG: sys.stderr.write(" all set.\n")
        self.seek(0)            # file pointer returns to the begining.
        return track

    def build_fwtrack (self, both_strand=False, do_merge=True):
        """Build FWTrackI from all lines, return a FWTrackI object.

        Note: All ranges will be merged (exclude the same
        range) then sorted after the track is built.

        If both_strand is True, it will store strand information in
        FWTrackI object.

        if do_merge is False, it will not merge the same range after
        the track is built.
        """
        fwtrack = FWTrackI()
        if LOG: sys.stdout.write(" adding ranges...\n")
        i = 0
        m = 0
        if LOG: sys.stdout.write("  "+prog_chars[i])
        for (chromosome,rg) in self:
            if not rg or not chromosome:
                continue
            if not both_strand:
                rg.strand = 1
            fwtrack.add_range(chromosome,rg)
            if LOG:
                i+=1
                if i == 50000:
                    m+=1
                    sys.stdout.write("\b"+str(m*50000)+"\n   ")
                    i=0
                sys.stdout.write("\b"+prog_chars[i-i/prog_num*prog_num])
            del rg
        if do_merge:
            if LOG: sys.stdout.write( "\b\bsorting then merging...\n")
            fwtrack.merge_overlap()
        if LOG: sys.stderr.write(" all set.\n")
        self.seek(0)            # file pointer returns to the begining.
        return fwtrack

    def next (self):
        (a,b) = self.readline()
        if a:
            return (a,b)
        else:
            raise StopIteration
        
    def readline (self, size=-1):
        """Read Range information from next line of the file, and return
        a tuple: (sequence_name, RangeI object).
        
        """
        thisline = file.readline(self, size)
        if not thisline: return (None,None)
        thisline = thisline.rstrip()
        if not thisline: return ("blank",None)
        if thisline[0] == "#": return ("comment line",None) # comment line is skipped
        thisfields = thisline.split(self.__delimiter__)
        if len(thisfields) > self.__strand_col__ :
            strand = thisfields[self.__strand_col__]
        else:
            strand = "1"        # default pos strand if no strand info can be found
        if strand == "+" or strand == "1":
            strand = 1
        elif strand == "-" or strand == "-1":
            strand = -1
        else:
            raise self.StrandFormatError(thisline,strand)

        #print thisfields[self.__chr_col__]
        return (thisfields[self.__chr_col__],
                RangeI(start  = int(thisfields[self.__start_col__])-self.__start_shift__,
                       end    = int(thisfields[self.__end_col__])-self.__start_shift__+self.__end_shift__,
                       strand = strand)
                )
        
    class FormatStringError(Exception):
        def __init__ (self, format):
            self.message = "Format string error: \"%s\"" % (format)

        def __str__ (self):
            return repr(self.message)

    class StrandFormatError(Exception):
        def __init__ (self, string, strand):
            self.message = "Strand information can not be recognized in this line: \"%s\",\"%s\"" % (string,strand)

        def __str__ (self):
            return repr(self.message)
